"""
评估
"""
import torch

from chatbot.chatbot_core.dataset import train_data_loader
from chatbot.chatbot_core.seq2seq import Seq2seq
import chatbot.config as config
from torch.optim import Adam
import torch.nn.functional as f
from tqdm import tqdm
from chatbot.lib import cut


def eval(by_word=True):
    seq2seq = Seq2seq().to(config.device)
    seq2seq.load_state_dict(torch.load(config.chatbot_model_save_path))
    optimizer = Adam(seq2seq.parameters(), lr=0.001)
    optimizer.load_state_dict(torch.load(config.chatbot_optimizer_save_path))
    while True:
        question = input("请输入：")
        question = cut(question, by_word=by_word)
        input_length = torch.LongTensor(
            [len(question) if len(question) > config.chatbot_input_max_len else config.chatbot_input_max_len]).to(
            config.device)
        question = torch.LongTensor(
            [config.chatbot_ws_input.transform(question, max_len=config.chatbot_input_max_len)]).to(config.device)
        indices = seq2seq.evaluation(question, input_length)
        output = "".join(config.chatbot_ws_target.inverse_transform(indices))
        print(output)
